import numpy as np
from sklearn.metrics import r2_score
from sklearn import datasets
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
# from sklearn.linear_model import LinearRegression

class LinearRegression:

    def __init__(self):
        self.coef = None
        self.interception = None
        self._theta = None
    def __repr__(self):
        return "LinearRegression"

    def fitNormal(self, xTrain, yTrain):
        assert xTrain.shape[0] == yTrain.shape[0], \
            'the size of xTrain must be equal to the size of yTrain'

        xB = np.hstack([np.ones((len(xTrain), 1)), xTrain])

        self._theta = np.linalg.inv(xB.T.dot(xB)).dot(xB.T).dot(yTrain)

        self.interception = self._theta[0]
        self.coef = self._theta[1:]

        return self

    def predict(self, xPredict):
        assert self.interception is not None and self.coef is not None, \
            'must fit befor predict'

        assert xPredict.shape[1] == len(self.coef), \
            'the feature number of xPredict must be equal to xTrain'

        xB = np.hstack([np.ones((len(xPredict), 1)), xPredict])

        return xB.dot(self._theta)

    def score(self, xTest, yTest):
        yPredict = self.predict(xTest)
        return r2_score(yTest, yPredict)

if __name__ == '__main__':
    linearRegression = LinearRegression()
    boston = datasets.load_boston()

    X = boston.data
    y = boston.target

    X = X[y < 50.0]
    y = y[y < 50.0]

    xTrain, xTest, yTrain, yTest = train_test_split(X, y, test_size=0.33, random_state=42)

    linearRegression.fitNormal(xTrain, yTrain)
    print(linearRegression.coef)
    print(linearRegression.score(xTest, yTest))